// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/operator.h>

namespace kumo::pollux::exec {
    class GroupId : public Operator {
    public:
        GroupId(
            int32_t operatorId,
            DriverCtx *driverCtx,
            const std::shared_ptr<const core::GroupIdNode> &groupIdNode);

        bool needsInput() const override;

        void addInput(RowVectorPtr input) override;

        RowVectorPtr getOutput() override;

        BlockingReason isBlocked(ContinueFuture * /*future*/) override {
            return BlockingReason::kNotBlocked;
        }

        bool isFinished() override {
            return finished_ || (noMoreInput_ && input_ == nullptr);
        }

    private:
        static constexpr column_index_t kMissingGroupingKey =
                std::numeric_limits<column_index_t>::max();

        bool finished_{false};

        /// A grouping set contains a subset of all the grouping keys. This list
        /// contains one entry per grouping set and identifies the grouping keys that
        /// are part of the set as indices of the input columns. The position in the
        /// list identifies the grouping key column in the output. Positions with
        /// kMissingGroupingKey correspond to grouping keys which are not included in
        /// the set.
        std::vector<std::vector<column_index_t> > groupingKeyMappings_;

        /// A list of input column indices corresponding to aggregation inputs. The
        /// position in the list identifies the column in the output.
        std::vector<column_index_t> aggregationInputs_;

        /// 'getOutput()' returns 'input_' for one grouping set at a time.
        /// 'groupingSetIndex_' contains the index of the grouping set to output in
        /// the next 'getOutput' call. This index is used to generate groupId column
        /// and lookup the input-to-output column mappings in the
        /// groupingKeyMappings_.
        int32_t groupingSetIndex_{0};
    };
} // namespace kumo::pollux::exec
